import pytesseract
import cv2
import os

"""
Tesseract-OCR 光学字符识别
"""


# 识别
def ocr(file_path):
    # 第一步，裁剪图片
    print('----------------第一步，裁剪图片--------------')
    img_cut_list = cut(file_path)
    for img_path in img_cut_list:
        print(img_path)
    print('-----------------裁剪图片完成-----------------')

    # 第二步，转换图片
    print('----------------第二步，转换图片--------------')
    img_convert_list = []
    for img_path in img_cut_list:
        img_convert_path = convert(img_path)
        img_convert_list.append(img_convert_path)

    for img_path in img_convert_list:
        print(img_path)
    print('-----------------转换图片全部完成-----------------')

    # 第三步，反色
    print('----------------第三步，反色图片--------------')
    img_inverse_list = []
    for img_path in img_convert_list:
        img_inverse_path = inverse(img_path)
        img_inverse_list.append(img_inverse_path)

    for img_path in img_inverse_list:
        print(img_path)
    print('----------------反色图片全部完成--------------')

    # 第四步，识别
    print('----------------第四步，识别图片--------------')
    img_dic = {}
    for img_path in img_inverse_list:
        result = tesseract(img_path)
        key = img_path.replace('\\convert\\inverse', '')
        img_dic[key] = result

    for key, value in img_dic.items():
        print(key)
        print(value)
    print('----------------识别图片全部完成--------------')

    return "success"


# 剪切图片，获取有用的信息，文件夹下需新建一个cut文件夹
def cut(file_path):
    import os.path as osp

    # 判断文件是否为图片
    def is_img(ext):
        ext.lower()
        if ext in ['.jpg', '.png', '.jpeg', '.bmp']:
            return True
        else:
            return False

    # 剪切图片
    def cut_img(file_path, img_name):
        img = cv2.imread(file_path + img_name)
        # print('高度、宽度、通道数:' + str(img.shape))  # 高度、宽度、通道数
        cropped = img[10:90, 0:650]  # 裁剪坐标为[y0:y1, x0:x1]
        new_img_path = file_path + "cut\\" + "cut_" + img_name
        cv2.imwrite(new_img_path, cropped)
        return new_img_path

    file_path = file_path + '\\'  # 文件路径
    new_img_list = []  # 存放新图片的地址
    # 返回指定的文件夹下包含的文件或文件夹名字的列表，这个列表按字母顺序排序。
    img_list = os.listdir(file_path)

    if not os.path.isdir(file_path + '\cut'):  # 创建cut文件夹
        os.makedirs(file_path + '\cut')

    for file in img_list:
        if is_img(osp.splitext(file)[1]):
            new_img_path = cut_img(file_path, file)
            new_img_list.append(new_img_path)

    return new_img_list


# 将图片转为黑底白字
def convert(img_path):
    from PIL import Image

    img_path_pre = img_path.split('.')[0]
    img_suffix = '.' + img_path.split('.')[1]  # 图片后缀
    img_name = img_path_pre.split('\\')[-1]  # 图片名
    file_path = img_path_pre.split(img_name)[0]  # 图片所在文件夹路径

    img = Image.open(img_path)

    # 先将图片灰度化
    # 模式L”为灰色图像，它的每个像素用8个bit表示，0表示黑，255表示白，其他数字表示不同的灰度。
    Img = img.convert('L')  # 灰度图

    # 自定义灰度界限，大于这个值为黑色，小于这个值为白色
    threshold = 235

    table = []
    for i in range(256):
        if i < threshold:
            table.append(0)
        else:
            table.append(1)

    if not os.path.isdir(file_path + '\\convert'):  # 创建convert文件夹
        os.makedirs(file_path + '\\convert')

    # 再将图片二值化
    photo = Img.point(table, '1')  # 二值化图
    img_convert_path = file_path + 'convert\\' + img_name + img_suffix
    photo.save(img_convert_path)
    return img_convert_path


# 图片反色，得到白底黑字，识别效果最好
def inverse(img_path):
    img_path_pre = img_path.split('.')[0]
    img_suffix = '.' + img_path.split('.')[1]  # 图片后缀
    img_name = img_path_pre.split('\\')[-1]  # 图片名
    file_path = img_path_pre.split(img_name)[0]  # 图片所在文件夹路径

    if not os.path.isdir(file_path + '\\inverse'):  # 创建convert文件夹
        os.makedirs(file_path + '\\inverse')

    src = cv2.imread(img_path, 1)
    Img = 255 - src
    img_inverse_path = file_path + 'inverse\\' + img_name + img_suffix
    cv2.imwrite(img_inverse_path, Img)
    return img_inverse_path


# 识别
def tesseract(img_path):
    # 字库
    lang = 'eng+chi_simp'
    # lang = 'eng+chi_sim'
    # lang = 'melang'
    result = pytesseract.image_to_string(img_path, lang)
    return result


if __name__ == '__main__':
    # file_path = r'C:\Users\dylan\Desktop\car'
    file_path = r'C:\Users\bingwa\Desktop\car'
    ocr(file_path)
